"""
This is rrt star code for 3D
@author: yue qi
"""
import time
import matplotlib.pyplot as plt
import numpy as np
from env3D import env
from utils3D import getDist, sampleFree, nearest, steer, isCollide, near, visualization, cost, path


class rrt():
    def __init__(self):
        self.env = env()
        self.Parent = {}
        self.V = []
        # self.E = edgeset()
        self.i = 0
        self.maxiter = 10000
        self.stepsize = 0.5
        self.Path = []
        self.done = False
        self.x0 = tuple(self.env.start)
        self.xt = tuple(self.env.goal)
        self.ind = 0
        # self.fig = plt.figure(figsize=(10, 8))

    def wireup(self, x, y):
        # self.E.add_edge([s, y])  # add edge
        self.Parent[x] = y

    def run(self):
        self.V.append(self.x0)
        plt.close('all')
        while self.ind < self.maxiter:
            xrand = sampleFree(self)
            xnearest = nearest(self, xrand)
            xnew, dist = steer(self, xnearest, xrand)
            collide, _ = isCollide(self, xnearest, xnew, dist=dist)
            if not collide:
                self.V.append(xnew)  # add point
                self.wireup(xnew, xnearest)

                if getDist(xnew, self.xt) <= self.stepsize:
                    self.wireup(self.xt, xnew)
                    self.Path, D = path(self)
                    # 三次多项式拟合平滑处理
                    degree = 3
                    smoothpath = np.array(self.Path)
                    smoothpath = smoothpath.reshape(-1, 3)
                    p_x = np.polyfit(np.arange(len(smoothpath[:, 0])), smoothpath[:, 0], degree)
                    p_y = np.polyfit(np.arange(len(smoothpath[:, 1])), smoothpath[:, 1], degree)
                    p_z = np.polyfit(np.arange(len(smoothpath[:, 2])), smoothpath[:, 2], degree)
                    # 生成拟合后的曲线
                    t = np.linspace(1, len(smoothpath[:, 0]), 50)
                    smoothed_x = np.polyval(p_x, t)
                    smoothed_y = np.polyval(p_y, t)
                    smoothed_z = np.polyval(p_z, t)
                    triplepath = np.vstack((smoothed_x, smoothed_y, smoothed_z)).T
                    triplepath = np.vstack((smoothpath[0, :], triplepath, smoothpath[-1, :]))

                    # 绘制原始路径和平滑后的路径
                    fig = plt.figure()
                    ax = fig.add_subplot(111, projection='3d')
                    ax.plot(smoothpath[:, 0], smoothpath[:, 1], smoothpath[:, 2], 'r', markersize=7,
                            markeredgecolor='r',
                            label='OriginPath')
                    ax.plot(triplepath[:, 0], triplepath[:, 1], triplepath[:, 2], 'b', markersize=7,
                            markeredgecolor='b',
                            label='SmoothedPath')
                    plt.xlabel('X(mm)')
                    plt.ylabel('Y(mm)')
                    plt.legend()
                    plt.show()
                    break
                visualization(self)
                self.i += 1
            self.ind += 1
            # if the goal is really reached
            
        self.done = True
        visualization(self)
        plt.xlabel('X(mm)')
        plt.ylabel('Y(mm)')
        plt.show()


if __name__ == '__main__':
    p = rrt()
    starttime = time.time()
    p.run()
    print('time used = ' + str(time.time() - starttime))
